import os
from typing import List, Optional
from pathlib import Path
from tqdm import tqdm
from llama_index.core import Document
from llama_index.core.readers import SimpleDirectoryReader
import magic.magic

class DocumentProcessor:
    """文档处理器类，用于处理和验证不同类型的文档"""
    
    SUPPORTED_MIME_TYPES = {
        'application/pdf': '.pdf',
        'text/plain': '.txt',
        'text/markdown': '.md',
        'text/html': '.html'
    }
    
    def __init__(self, documents_dir: str = "./documents"):
        self.documents_dir = Path(documents_dir)
        self.documents_dir.mkdir(exist_ok=True)
        
    def validate_file(self, file_path: Path) -> bool:
        """验证文件是否支持"""
        try:
            mime_type = magic.magic.from_file(str(file_path), mime=True)
            return mime_type in self.SUPPORTED_MIME_TYPES
        except Exception as e:
            print(f"验证文件时出错 {file_path}: {str(e)}")
            return False
    
    def process_documents(self) -> List[Document]:
        """处理文档目录中的所有文件"""
        documents = []
        
        # 获取所有文件
        files = list(self.documents_dir.glob("**/*"))
        valid_files = []
        
        # 验证文件
        print("正在验证文件...")
        for file_path in tqdm(files):
            if file_path.is_file() and self.validate_file(file_path):
                valid_files.append(file_path)
            elif file_path.is_file():
                print(f"警告: 不支持的文件类型 {file_path}")
        
        if not valid_files:
            print(f"在 {self.documents_dir} 中没有找到有效文档")
            return documents
        
        # 读取文档
        print("正在处理文档...")
        reader = SimpleDirectoryReader(input_files=[str(f) for f in valid_files])
        documents = reader.load_data()
        
        print(f"成功处理了 {len(documents)} 个文档")
        return documents
    
    def get_document_stats(self) -> dict:
        """获取文档统计信息"""
        stats = {
            "total_files": 0,
            "valid_files": 0,
            "file_types": {},
            "total_size": 0
        }
        
        for file_path in self.documents_dir.glob("**/*"):
            if file_path.is_file():
                stats["total_files"] += 1
                stats["total_size"] += file_path.stat().st_size
                
                if self.validate_file(file_path):
                    stats["valid_files"] += 1
                    ext = file_path.suffix
                    stats["file_types"][ext] = stats["file_types"].get(ext, 0) + 1
        
        return stats

    @staticmethod
    def get_document_files(directory: str) -> List[Path]:
        """
        获取目录中的所有文档文件
        
        Args:
            directory (str): 文档目录路径
            
        Returns:
            List[Path]: 文档文件路径列表
        """
        supported_extensions = {'.txt', '.md', '.pdf', '.html', '.doc', '.docx'}
        document_files = []
        
        for root, _, files in os.walk(directory):
            for file in files:
                if Path(file).suffix.lower() in supported_extensions:
                    document_files.append(Path(root) / file)
                    
        return document_files
    
    @staticmethod
    def validate_document(file_path: Path) -> bool:
        """
        验证文档是否有效
        
        Args:
            file_path (Path): 文档文件路径
            
        Returns:
            bool: 文档是否有效
        """
        if not file_path.exists():
            return False
            
        # 检查文件大小（限制为 10MB）
        if file_path.stat().st_size > 10 * 1024 * 1024:
            return False
            
        return True
    
    @staticmethod
    def preprocess_text(text: str) -> str:
        """
        预处理文本内容
        
        Args:
            text (str): 原始文本
            
        Returns:
            str: 处理后的文本
        """
        # 移除多余的空白字符
        text = ' '.join(text.split())
        # 这里可以添加更多的文本预处理逻辑
        return text